

import torch
import torch.nn as nn
import torch.nn.functional as F

class Generator(nn.Module):

    def __init__(self, state_dim, action_dim):
        super(Generator, self).__init__()

        self.fc1 = nn.Linear(state_dim + action_dim + 1, 128)
        self.fc2 = nn.Linear(128, 64)
        self.fc3 = nn.Linear(64, 1)

    def forward(self, state, action, noise):
        state_action = torch.cat([state, action, noise], 1)

        q = F.relu(self.fc1(state_action))
        q = F.relu(self.fc2(q))
        q = self.fc3(q)
        return q


class Generator_target(nn.Module):

    def __init__(self, state_dim, action_dim):
        super(Generator_target, self).__init__()

        self.fc1 = nn.Linear(state_dim + action_dim + 1, 128)
        self.fc2 = nn.Linear(128, 64)
        self.fc3 = nn.Linear(64, 1)

    def forward(self, state, action, noise):
        state_action = torch.cat([state, action, noise], 1)

        q = F.relu(self.fc1(state_action))
        q = F.relu(self.fc2(q))
        q = self.fc3(q)
        return q


class Discriminator(nn.Module):

    def __init__(self, ):
        super(Discriminator, self).__init__()

        self.fc1 = nn.Linear(1, 128)
        self.fc2 = nn.Linear(128, 64)
        self.fc3 = nn.Linear(64, 1)

    def forward(self, action_value):
        # state_action = torch.cat([action_value, action_value_target], 1)
        state_action = action_value

        q = F.relu(self.fc1(state_action))
        q = F.relu(self.fc2(q))
        q = self.fc3(q)
        return q